'''
用于生成annotation.txt文件，格式为：image_path xmin,ymin,xmax,ymax,class_ind xmin,ymin,xmax,ymax,class_ind
仅用于argusswift的YOLOv4-pytorch实现及其衍生算法实现，如RISFNet
'''

import os
from tqdm.contrib import tzip
import xml.etree.ElementTree as ET
import config.risfnet_config as cfg

def parse_flow_annotation(data_path, file_type, anno_path, classes, use_difficult_bbox=False):
    """
    phase pascal flow annotation, eg:[image_global_path xmin,ymin,xmax,ymax,cls_id]
    :param data_path: eg: FloW/FloW-RI
    :param file_type: eg: 'trainval''train''val'
    :param anno_path: path to ann file
    :param use_difficult_bbox: whither use different sample
    :return: batch size of data set
    """

    with open(os.path.join(data_path, file_type+'.txt'), 'r') as f:
        lines = [x.strip() for x in f.readlines()]
        lines = [x.split(' ') for x in lines]
    image_files = [lin[0] for lin in lines]
    label_files = [lin[1] for lin in lines]

    with open(anno_path, 'a') as f:
        for image_file, label_file in tzip(image_files, label_files):
            annotation = image_file
            label_path = os.path.join(data_path, label_file)
            root = ET.parse(label_path).getroot()
            objects = root.findall('object')
            for obj in objects:
                # difficult = obj.find("difficult").text.strip() # original code
                difficult = obj.find("difficult").text.strip() if obj.find("difficult") is not None else obj.find("Difficult").text.strip() # @@identify different keys
                if (not use_difficult_bbox) and (int(difficult) == 1): continue # difficult 表示是否容易识别，0表示容易，1表示困难
                bbox = obj.find('bndbox')
                class_id = classes.index(obj.find("name").text.lower().strip())
                xmin = bbox.find('xmin').text.strip()
                ymin = bbox.find('ymin').text.strip()
                xmax = bbox.find('xmax').text.strip()
                ymax = bbox.find('ymax').text.strip()
                annotation += ' ' + ','.join([xmin, ymin, xmax, ymax, str(class_id)])
            annotation += '\n'
            # print(annotation)
            f.write(annotation)
    return len(image_files)

if __name__ =="__main__":
    CLASSES = cfg.Customer_DATA["CLASSES"] # dataset class
    DATA_PATH = cfg.DATA_PATH

    # trainval_set
    train_annotation_path = os.path.join(DATA_PATH, 'train_annotation.txt')
    if os.path.exists(train_annotation_path):
        os.remove(train_annotation_path)

    # test_set
    test_annotation_path = os.path.join(DATA_PATH, 'test_annotation.txt')
    if os.path.exists(test_annotation_path):
        os.remove(test_annotation_path)

    len_train = parse_flow_annotation(DATA_PATH, "trainval", train_annotation_path, CLASSES) # 从trainval.txt提取注释
    len_test = parse_flow_annotation(DATA_PATH, "test", test_annotation_path, CLASSES) # 从test.txt提取注释

    print("The number of images for train and test are :train : {0} | test : {1}".format(len_train, len_test))
